
/*
* Copyright (c) 2008 - 2009 NVIDIA Corporation.  All rights reserved.
*
* NVIDIA Corporation and its licensors retain all intellectual property and proprietary
* rights in and to this software, related documentation and any modifications thereto.
* Any use, reproduction, disclosure or distribution of this software and related
* documentation without an express license agreement from NVIDIA Corporation is strictly
* prohibited.
*
* TO THE MAXIMUM EXTENT PERMITTED BY APPLICABLE LAW, THIS SOFTWARE IS PROVIDED *AS IS*
* AND NVIDIA AND ITS SUPPLIERS DISCLAIM ALL WARRANTIES, EITHER EXPRESS OR IMPLIED,
* INCLUDING, BUT NOT LIMITED TO, IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
* PARTICULAR PURPOSE.  IN NO EVENT SHALL NVIDIA OR ITS SUPPLIERS BE LIABLE FOR ANY
* SPECIAL, INCIDENTAL, INDIRECT, OR CONSEQUENTIAL DAMAGES WHATSOEVER (INCLUDING, WITHOUT
* LIMITATION, DAMAGES FOR LOSS OF BUSINESS PROFITS, BUSINESS INTERRUPTION, LOSS OF
* BUSINESS INFORMATION, OR ANY OTHER PECUNIARY LOSS) ARISING OUT OF THE USE OF OR
* INABILITY TO USE THIS SOFTWARE, EVEN IF NVIDIA HAS BEEN ADVISED OF THE POSSIBILITY OF
* SUCH DAMAGES
*/

#include <optix.h>
#include <optixu/optixu_math_namespace.h>
#include "helpers.h"
#include "spectral_path_tracer.h"
#include "random.h"

using namespace optix;


// Scene wide
rtDeclareVariable(float, scene_epsilon, , );
rtDeclareVariable(rtObject, top_object, , );

// For camera
rtDeclareVariable(float3, eye, , );
rtDeclareVariable(float3, U, , );
rtDeclareVariable(float3, V, , );
rtDeclareVariable(float3, W, , );
rtDeclareVariable(float3, bad_color, , );
rtDeclareVariable(unsigned int, frame_number, , );
rtDeclareVariable(unsigned int, sqrt_num_samples, , );

rtBuffer<TriangleLight>     lights;

//ray types
rtDeclareVariable(unsigned int, pathtrace_ray_type, , );
rtDeclareVariable(unsigned int, pathtrace_shadow_ray_type, , );
rtDeclareVariable(unsigned int, rr_begin_depth, , );
rtDeclareVariable(unsigned int, max_depth, , );
rtDeclareVariable(int, lgt_idx, attribute lgt_idx, );
rtDeclareVariable(float, Y_log_av, , );
rtDeclareVariable(float, Y_max, , );

rtDeclareVariable(float, lambdaStart, , );
rtDeclareVariable(float, lambdaEnd, , );
rtDeclareVariable(float, lambdaDelta, , );


//output buffers
rtBuffer<float4, 2>              output_buffer;
rtBuffer<float3, 2>              hdr_buffer;

//rays
rtDeclareVariable(PerRayData_pathtrace, current_prd, rtPayload, );

//optix tracked data
rtDeclareVariable(optix::Ray, ray, rtCurrentRay, );
rtDeclareVariable(float, t_hit, rtIntersectionDistance, );
rtDeclareVariable(uint2, launch_index, rtLaunchIndex, );
rtDeclareVariable(uint2, launch_dim, rtLaunchDim, );

//geometry attributes
rtDeclareVariable(float3, geometric_normal, attribute geometric_normal, );
rtDeclareVariable(float3, shading_normal, attribute shading_normal, );
rtDeclareVariable(float3, texcoord, attribute texcoord, );


// For miss program
rtDeclareVariable(float3, bg_color, , );

//-----------------------------------------------------------------------------
//
//  Camera program -- main ray tracing loop
//
//-----------------------------------------------------------------------------

RT_PROGRAM void pathtrace_camera()
{
	size_t2 screen = output_buffer.size();

	float2 inv_screen = 1.0f / make_float2(screen) * 2.f;
	float2 pixel = (make_float2(launch_index)) * inv_screen - 1.f;

	float2 jitter_scale = inv_screen / sqrt_num_samples;
	unsigned int samples_per_pixel = sqrt_num_samples*sqrt_num_samples;
	float3 result = make_float3(0.0f);
	float spd[16];
	unsigned int seed = tea<16>(screen.x*launch_index.y + launch_index.x, frame_number);

	do {
		unsigned int x = samples_per_pixel%sqrt_num_samples;
		unsigned int y = samples_per_pixel / sqrt_num_samples;
		float2 jitter = make_float2(x - rnd(seed), y - rnd(seed));
		float2 d = pixel + jitter*jitter_scale;
		float3 ray_origin = eye;
		float3 ray_direction = normalize(d.x*U + d.y*V + W);

		PerRayData_pathtrace prd;
		prd.result = make_float3(0.f);
		float lambda = lambdaStart;

		prd.attenuation = 1.f;
		prd.countEmitted = true;
		prd.done = false;
		prd.inside = false;
		prd.seed = seed;
		prd.depth = 0;
		for (int sspx = 0; sspx < 16; sspx++) {
			prd.lambda = lambda;
			for (;;)
			{
				Ray ray = make_Ray(ray_origin, ray_direction, pathtrace_ray_type, scene_epsilon, RT_DEFAULT_MAX);
				rtTrace(top_object, ray, prd);
				if (prd.done) {
					spd[sspx] += prd.radiance * prd.attenuation;
					break;
				}

				// RR
				if (prd.depth >= rr_begin_depth){
					//break;
					float pcont = prd.attenuation;
					if (rnd(prd.seed) >= pcont || prd.depth > max_depth)
						break;
					prd.attenuation /= pcont;
				}
				prd.depth++;
				spd[sspx] += prd.radiance * prd.attenuation;
				ray_origin = prd.origin;
				ray_direction = prd.direction;
			}
			lambda += lambdaDelta;
		} // eye ray


		float3 pix;
		XYZToRGB(SPD2XYZ(create_spd(lambdaStart, lambdaEnd, spd)), pix);
		result += pix;
		seed = prd.seed;
	} while (--samples_per_pixel);

	float3 pixel_color = result / (sqrt_num_samples*sqrt_num_samples);

	if (frame_number > 1)
	{
		float a = 1.0f / (float)frame_number;
		float b = ((float)frame_number - 1.0f) * a;
		hdr_buffer[launch_index] = a * pixel_color + b * hdr_buffer[launch_index];
	}
	else
	{
		hdr_buffer[launch_index] = pixel_color;
	}
	output_buffer[launch_index] = make_float4(hdr_buffer[launch_index], 0.0f);
}


rtDeclareVariable(float, time_view_scale, , ) = 1e-6f;

rtDeclareVariable(float, aperture_radius, , );
rtDeclareVariable(float, focal_scale, , );


rtDeclareVariable(SPD16, emission_color, , );


rtDeclareVariable(SPD16, reflectance, , );
//rtDeclareVariable(PerRayData_pathtrace_bsdf_shadow, bsdf_shadow_prd, rtPayload, );

rtDeclareVariable(float3, directional_light, , );
rtDeclareVariable(float3, directional_light_col, , );
rtDeclareVariable(float, solid_angle, , ) = 6.0e-5f * 100.0f;
rtDeclareVariable(float, phong_exp, , );
rtDeclareVariable(float3, glass_color, , );



RT_PROGRAM void diffuse()
{


	/*if (length(emission_color) > 0.0f)
	{
		current_prd.radiance = current_prd.countEmitted ? emission_color : 1.f;
		current_prd.done = true;
		return;
	}*/

	float3 world_shading_normal = normalize(rtTransformNormal(RT_OBJECT_TO_WORLD, shading_normal));
	float3 world_geometric_normal = normalize(rtTransformNormal(RT_OBJECT_TO_WORLD, geometric_normal));
	float2 uv = make_float2(texcoord);
	float Kd = sample_spd(reflectance, current_prd.lambda);
	float3 ffnormal = faceforward(world_shading_normal, -ray.direction, world_geometric_normal);

	float3 hitpoint = ray.origin + t_hit * ray.direction;
	current_prd.origin = hitpoint;

	float z1 = rnd(current_prd.seed);
	float z2 = rnd(current_prd.seed);
	float3 p;
	cosine_sample_hemisphere(z1, z2, p);
	float3 v1, v2;
	createONB(ffnormal, v1, v2);
	//ffnormal = normalize(make_float3( v1.x * mappedNormal.x + v2.x * mappedNormal.y + ffnormal.x * mappedNormal.z, v1.y * mappedNormal.x + v2.y * mappedNormal.y + ffnormal.y * mappedNormal.z, v1.z * mappedNormal.x + v2.z * mappedNormal.y + ffnormal.z * mappedNormal.z));
	current_prd.direction = v1 * p.x + v2 * p.y + ffnormal * p.z;
	current_prd.attenuation = current_prd.attenuation * Kd; // use the diffuse_color as the diffuse response
	current_prd.countEmitted = false;

	unsigned int num_lights = lights.size();
	float result = 0.f;

	//for(int i = 0; i < num_lights; ++i) 
	{
		TriangleLight light = lights[rnd(current_prd.seed)*(num_lights)];
		//TriangleLight light = lights[i];

		float z1 = rnd(current_prd.seed);
		float z2 = rnd(current_prd.seed);


		float alpha = 1.0f - sqrt(fold(z1));
		float beta = (1.0f - fold(z2)) * sqrt(fold(z1));
		float gamma = fold(z2) * sqrt(fold(z1));
		float3 light_pos = light.v3 * gamma + light.v1 * alpha + light.v2 * beta;

		float Ldist = length(light_pos - hitpoint);
		float3 L = normalize(light_pos - hitpoint);
		float nDl = dot(ffnormal, L);
		float LnDl = dot(light.normal, L);
		float A = length(cross(light.v2 - light.v3, light.v1 - light.v3));

		// cast shadow ray
		if (nDl > 0.0f && LnDl > 0.0f)
		{
			PerRayData_pathtrace_shadow shadow_prd;
			shadow_prd.inShadow = false;
			Ray shadow_ray = make_Ray(hitpoint, L, pathtrace_shadow_ray_type, scene_epsilon, Ldist);
			rtTrace(top_object, shadow_ray, shadow_prd);

			if (!shadow_prd.inShadow)
			{
				float weight = nDl * LnDl / (M_PIf*Ldist*Ldist) * A;
				float le = sample_spd(light.emission, current_prd.lambda);
				result += num_lights * le* weight;
			}
		}
	}
	current_prd.radiance = result;
	//current_prd.result += current_prd.attenuation * result;
	//current_prd.radiance = make_float3(0.0f);
	//current_prd.attenuation = current_prd.attenuation * Kd; 
}


rtDeclareVariable(PerRayData_pathtrace_shadow, current_prd_shadow, rtPayload, );

RT_PROGRAM void shadow()
{
	if (emission_color.delta == 0.0f)
	{
		current_prd_shadow.inShadow = true;
		current_prd_shadow.attenuation = 0.0f;
		rtTerminateRay();
	}
}



//-----------------------------------------------------------------------------
//
//  Exception program
//
//-----------------------------------------------------------------------------

RT_PROGRAM void exception()
{
	output_buffer[launch_index] = make_float4(bad_color, 0.0f);
}


//-----------------------------------------------------------------------------
//
//  Miss program
//
//-----------------------------------------------------------------------------

RT_PROGRAM void miss()
{
	current_prd.radiance = 0.f;
	current_prd.done = true;
}
